import codecs
import numpy as np
from gensim import corpora
from gensim.models import LdaModel
import matplotlib.pyplot as plt

# 读取数据
train = []
fp = codecs.open('CNKI-output.txt', 'r', encoding='utf8')
for line in fp:
    if line != '':
        line = line.split()
        train.append([w for w in line])

# 创建词典和语料库
dictionary = corpora.Dictionary(train)
corpus = [dictionary.doc2bow(text) for text in train]

# 测试不同主题数
topic_numbers = range(1, 16)  # 从1到15
perplexities = []

# 计算每个主题数下的困惑度
for num_topics in topic_numbers:
    lda = LdaModel(corpus=corpus, 
                   id2word=dictionary, 
                   num_topics=num_topics, 
                   passes=60)
    perplexity = lda.log_perplexity(corpus)
    perplexities.append(perplexity)
    print(f"Topics: {num_topics}, Perplexity: {perplexity}")

# 绘制困惑度曲线
plt.figure(figsize=(10, 6))
plt.plot(topic_numbers, perplexities, 'b-', marker='o')
plt.xlabel('Number of Topics')
plt.ylabel('Log Perplexity')
plt.title('LDA Model Perplexity vs Number of Topics')
plt.grid(True)
plt.savefig('topicnum/perplexity_plot.png')
plt.close()

# 找出局部最小值
local_mins = []
for i in range(1, len(perplexities)-1):
    if perplexities[i] < perplexities[i-1] and perplexities[i] < perplexities[i+1]:
        local_mins.append((topic_numbers[i], perplexities[i]))

# 将结果写入文件
with open('topicnum/results.txt', 'w', encoding='utf-8') as f:
    f.write("困惑度局部最小值出现在以下主题数：\n")
    for topic_num, perplexity in local_mins:
        f.write(f"主题数: {topic_num}, 困惑度: {perplexity}\n")
    
    # 找出全局最小值
    min_perplexity_idx = np.argmin(perplexities)
    f.write(f"\n全局最小值：\n")
    f.write(f"主题数: {topic_numbers[min_perplexity_idx]}, 困惑度: {perplexities[min_perplexity_idx]}") 